import librosa
import numpy as np
import keras
from keras.models import Sequential
from keras.layers import Dense
import math
# from keras_radam import RAdam
# import lookahead

print('读取训练信号\n')
for x in range(500):
    y, sr = librosa.load(r'C:\Users\WangCan\PycharmProjects\speech\venv\datasource\s16\S16_'+str(x+1)+'.wav', sr=None)
    y = y.reshape(1, len(y))
    if(x == 0):
        femaleSignal = y
    else:
        femaleSignal = np.c_[femaleSignal, y]
    y, sr = librosa.load(r'C:\Users\WangCan\PycharmProjects\speech\venv\datasource\s22\S22_' + str(x + 1) + '.wav', sr=None)
    y = y.reshape(1, len(y))
    if (x == 0):
        female2Signal = y
    else:
        female2Signal = np.c_[female2Signal, y]
'#21256500 21695500'
print(len(femaleSignal[0]), len(female2Signal[0]))
femaleSignal = np.array(femaleSignal[0][0:21256192]).reshape(1, -1)
female2Signal = np.array(female2Signal[0][0:21256192]).reshape(1, -1)
print(len(femaleSignal[0]))

print("读取测试信号\n")
for x in range(10):
    y, sr = librosa.load(r'C:\Users\WangCan\PycharmProjects\speech\venv\datasource\s16\S16_'+str(x+701)+'.wav', sr=None)
    y = y.reshape(1, len(y))
    if(x == 0):
        cFemaleSignal = y
    else:
        cFemaleSignal = np.c_[cFemaleSignal, y]
    y, sr = librosa.load(r'C:\Users\WangCan\PycharmProjects\speech\venv\datasource\s22\S22_' + str(x + 601) + '.wav', sr=None)
    y = y.reshape(1, len(y))
    if (x == 0):
        cFemale2Signal = y
    else:
        cFemale2Signal = np.c_[cFemale2Signal, y]
'#365010 517760'
print(len(cFemaleSignal[0]), len(cFemale2Signal[0]))
cFemaleSignal = np.array(cFemaleSignal[0][0:423424]).reshape(1, -1)
cFemale2Signal = np.array(cFemale2Signal[0][0:423424]).reshape(1, -1)
print(len(cFemaleSignal[0]), len(cFemaleSignal[0]) == len(cFemale2Signal[0]))
cMixedSignal = np.array(cFemaleSignal[0]+cFemale2Signal[0]).reshape(1, -1)


print("信号变换：")
#训练女生2信号
bFemale2Signal = librosa.stft(female2Signal[0], n_fft=512)
bFemale2SignalAmplitude = np.abs(bFemale2Signal)
bFemale2SignalAngle = np.angle(bFemale2Signal)
#训练女生
bFemaleSignal = librosa.stft(femaleSignal[0], n_fft=512)
bFemaleSignalAmplitude = np.abs(bFemaleSignal)
bFemaleSignalAngle = np.angle(bFemaleSignal)
#训练  目标IRM
IRM = bFemaleSignalAmplitude/(bFemaleSignalAmplitude+bFemale2SignalAmplitude)
TrainIRM = IRM.transpose()
#训练输入混合信号
bMixedSignalAmplitude = bFemaleSignalAmplitude+bFemale2SignalAmplitude
bMixedSignalMax = np.max(bMixedSignalAmplitude, axis=0)
bMixedSignalNorm = bMixedSignalAmplitude/bMixedSignalMax
TrainMixedSignal = bMixedSignalNorm.transpose()
#测试混合信号
tMixedSignal = librosa.stft(cMixedSignal[0], n_fft=512)
tMixedSignalAmplitude = np.abs(tMixedSignal)
tMixedSignalAngle = np.angle(tMixedSignal)
tMixedSignalAmplitudeMax = np.max(tMixedSignalAmplitude, axis=0)
tMixedSignalAmplitudeNorm = tMixedSignalAmplitude/tMixedSignalAmplitudeMax
TestMixedSignal = tMixedSignalAmplitudeNorm.transpose()
print("信号变换完成")

print("训练开始：")
X_Train = TrainMixedSignal
Y_Train = TrainIRM
X_Test = TestMixedSignal
model = Sequential()
model.add(Dense(1024, activation='relu', input_dim=257))
model.add(Dense(1024, activation='selu'))
model.add(Dense(1024, activation='sigmoid'))
model.add(Dense(257, activation='relu'))
sgd = keras.optimizers.SGD(lr=0.01, momentum=0.9)
NADAM =keras.optimizers.Nadam(lr=0.002, beta_1=0.9, beta_2=0.999)
model.compile(optimizer="adam", loss='mse', metrics=['mae'])
# lookahead1 = lookahead.Lookahead(k=5, alpha=0.5)
# lookahead1.inject(model)
model.fit(X_Train, Y_Train, epochs=200, batch_size=256, verbose=2)
Y_pred = model.predict(X_Test, batch_size=256, verbose=0, steps=None)
print("训练完成;")

print("语音输出：")
rFemaleSignalAmplitude = (Y_pred*TestMixedSignal).transpose()*tMixedSignalAmplitudeMax
rFemaleSignal = rFemaleSignalAmplitude*np.power(math.e, 1j*tMixedSignalAngle)
rebuildFemaleSignal = librosa.istft(rFemaleSignal)
librosa.output.write_wav(r".\speech\f2\IRM\adam\rebuild_signal.wav", rebuildFemaleSignal, sr=25000)
librosa.output.write_wav(r".\speech\f2\IRM\adam\origin_signal.wav", cFemaleSignal[0], sr=25000)
# librosa.output.write_wav(r".\speech\IRM\ranger\rebuild_signal.wav", rebuildMaleSignal, sr=25000)
# librosa.output.write_wav(r".\speech\IRM\ranger\origin_signal.wav", cMixedSignal[0], sr=25000)
print("语音输出完成")
